#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Created by Ross on 19-4-14, contributor: Ross, GabrielThompson

# bert server启动命令
# bert-serving-start -model_dir chinese_L-12_H-768_A-12 -max_seq_len 32
"""
用法： python3 gen_bert_emb.py src_path result_path
其中 src_path 是一个文件夹路径，里面存放着数据文件，数据文件里面有多行，没行一个句子（已经过分词）
"""


import os
import sys

import numpy as np
from bert_serving.client import BertClient

BERT_MAX_SEQ_LEN = 30
bc = BertClient()


def loaddata(src_path):
    data_path = src_path  # '/home/gabriel/文档/10fold_shijie_segment'   存放你要转换词向量的十折数据集文件夹
    name_list = os.listdir(data_path)
    files = []
    data_list = []
    for name in name_list:
        files.append(os.path.join(data_path, name))

    print('data loaded')
    for file in files:
        tmp = []
        data = open(file, 'r', encoding='utf-8')
        for item in data:
            # print(item)
            sentence = item.strip()
            tmp.append(sentence)

        data_list.append(tmp)
    print('append finished')
    return data_list, name_list


def store_Bert_npy(bert_vec, length, name):
    for sent_bert, true_len in zip(bert_vec, length):
        assert len(sent_bert) == true_len
    np.save(os.path.join(result_path, str(name).replace('.txt', '') + '.npy'), bert_vec)


def to_Bert_npy(data_list, name_list):
    bert_vec = []
    cnt = 0
    for data, name in zip(data_list, name_list):
        cnt += 1
        lengths = [len(x.split()) for x in data]  # 句子长度

        new_data = []
        _data = [s.split() for s in data]
        for sentence in _data:
            new_data.append(bc.encode(sentence))

        store_Bert_npy(new_data, lengths, name)
        print(np.shape(new_data))
        print('file' + str(cnt) + 'finished')
    return bert_vec


if __name__ == '__main__':
    src_path = sys.argv[1]  # 第一个参数
    result_path = sys.argv[2]  # 第二个参数
    print('src_path:', src_path)
    print('result_path:', result_path)
    if not os.path.exists(result_path):
        os.makedirs(result_path)

    data_list, name_list = loaddata(src_path)
    bert_vec = to_Bert_npy(data_list, name_list)
    print('finished')
